package com.tamude.gpt.service;

import com.alibaba.excel.EasyExcel;
import com.tamude.gpt.EmbeddingQASystem;
import com.tamude.gpt.TestEntity;

import java.io.IOException;
import java.io.InputStream;
import java.util.*;
import java.util.stream.Collectors;

public class SearchFunction {

    public static EmbeddingQASystem qaSystem = new EmbeddingQASystem();

    // 定义Embedding模型的调用方式
    private static List<Double> getEmbedding(String query) throws IOException {
        // 调用你的embedding API来获取embedding
        return qaSystem.getEmbedding(query);
    }

    // 计算余弦相似度
    private static double cosineSimilarity(List<Double> vecA, List<Double> vecB) {
        double dotProduct = 0.0;
        double normA = 0.0;
        double normB = 0.0;

        for (int i = 0; i < vecA.size(); i++) {
            dotProduct += vecA.get(i) * vecB.get(i);
            normA += Math.pow(vecA.get(i), 2);
            normB += Math.pow(vecB.get(i), 2);
        }

        normA = Math.sqrt(normA);
        normB = Math.sqrt(normB);

        return (normA == 0 || normB == 0) ? 0 : dotProduct / (normA * normB);
    }

    // 返回最相关的字符串及其相关性
    public static List<Map.Entry<String, Double>> stringsRankedByRelatedness(
            String query, List<Map<String, Object>> df, int topN) throws IOException {

        List<Double> queryEmbedding = getEmbedding(query);

        // 将字符串与相关性配对
        List<Map.Entry<String, Double>> stringsAndRelatednesses = new ArrayList<>();
        for (Map<String, Object> row : df) {
            List<Double> rowEmbedding = (List<Double>) row.get("embedding");
            String text = (String) row.get("text");
            double relatedness = cosineSimilarity(queryEmbedding, rowEmbedding);
            stringsAndRelatednesses.add(new AbstractMap.SimpleEntry<>(query, relatedness));
        }

        // 按相关性从高到低排序
        stringsAndRelatednesses = stringsAndRelatednesses.stream()
                .sorted((a, b) -> Double.compare(b.getValue(), a.getValue()))
                .collect(Collectors.toList());

        // 返回前N个相关的字符串及其相关性
        return stringsAndRelatednesses.subList(0, Math.min(topN, stringsAndRelatednesses.size()));
    }

    public static void main(String[] args) throws IOException {
        // 示例数据
        List<Map<String, Object>> df = new ArrayList<>();
        Map<String, Object> row1 = new HashMap<>();
        String question = "游戏宅在什么时间购买了什么商品？他对该商品的评价如何呢？";
        row1.put("text", question);
        row1.put("embedding", qaSystem.getEmbedding("游戏宅在什么时间购买了什么商品？他对该商品的评价如何呢？"));
        df.add(row1);

        // 调用函数
        SearchFunction searchFunction = new SearchFunction();
        List<String> readExcel = searchFunction.readExcel();
        List<Map.Entry<String, Double>> results = new ArrayList<>();
        for (String q : readExcel) {
            try {
                results.addAll(stringsRankedByRelatedness(q,df,100));
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }

        results = results.stream().sorted(Map.Entry.comparingByValue()).collect(Collectors.toList());
        Collections.reverse(results);
        if (results.size() > 5) {
            results = results.subList(0, 5);
        }
        StringBuilder sb = new StringBuilder();
        sb.append("基于我提供的下列内容：\n");
        results.forEach(r -> sb.append(r.getKey()));
        sb.append("\n回答我的问题 ：\n")
                .append(question);
        System.out.println(sb);
//        System.out.println(qaSystem.generateAnswer(sb.toString()));
    }


    public List<String> readExcel(){
        InputStream inputStream = this.getClass().getClassLoader().getResourceAsStream("embedding测试数据.xlsx");
        List<TestEntity> list = EasyExcel.read(inputStream)
                .head(TestEntity.class)
                // 设置sheet,默认读取第一个
                .sheet()
                // 设置标题所在行数
                .headRowNumber(0)
                .doReadSync();

        List<String> results = new ArrayList<>();
        TestEntity headEntity = list.get(0);
        for (int i = 1; i < list.size(); i++) {
            StringBuilder sb = new StringBuilder();
            TestEntity e = list.get(i);
            sb.append(headEntity.getUser()).append(" | ").append(headEntity.getCommodity()).append(" | ").append(headEntity.getBuyDate()).append(" | ").append(headEntity.getStars()).append(" | ").append(headEntity.getEvaluation()).append(" | ").append(headEntity.getEvaluationDate()).append("\n");
            sb.append(e.getUser()).append(" | ").append(e.getCommodity()).append(" | ").append(e.getBuyDate()).append(" | ").append(e.getStars()).append(" | ").append(e.getEvaluation()).append(" | ").append(e.getEvaluationDate()).append("\n");
            results.add(sb.toString());
        }

        return results;
    }
}
